package com.guchenbo.lssp.pool;

import cn.hutool.core.net.NetUtil;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.guchenbo.lssp.core.Session;
import com.guchenbo.lssp.core.Statement;
import com.guchenbo.lssp.domain.*;
import com.guchenbo.lssp.exception.SessionException;
import com.guchenbo.lssp.exception.SessionNotReadyException;
import com.guchenbo.lssp.exception.StatementNotReadyException;
import com.guchenbo.lssp.support.LivyProperties;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.retry.backoff.FixedBackOffPolicy;
import org.springframework.retry.policy.SimpleRetryPolicy;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.web.client.HttpClientErrorException;
import org.springframework.web.client.HttpServerErrorException;
import org.springframework.web.client.RestTemplate;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * @author guchenbo
 */
public class LivyHelper {
    protected static final Logger log = LoggerFactory.getLogger(LivyHelper.class);

    public static final String REGEX = "(.*)\\.(.*)\\((.*):(\\d+)\\)";
    public static final Pattern PATTERN = Pattern.compile(REGEX);
    public static final String SESSION_PREFIX = "LSSP-%s-";

    public static final String SESSION_URL = "%s/sessions";
    public static final String SESSION_ID_URL = "%s/sessions/%s";
    public static final String STATEMENT_URL = "%s/sessions/%s/statements";
    public static final String STATEMENT_ID_URL = "%s/sessions/%s/statements/%s";

    private final String livyHost;
    private final RestTemplate restTemplate;
    private final LivyProperties livyProperties;
    private RetryTemplate sessionRetryTemplate;
    private RetryTemplate statementRetryTemplate;

    public LivyHelper(String livyHost, RestTemplate restTemplate, LivyProperties livyProperties) {
        if (livyHost != null && livyHost.endsWith("/")) {
            livyHost = livyHost.substring(0, livyHost.length() - 1);
        }
        this.livyHost = livyHost;
        this.restTemplate = restTemplate;
        this.livyProperties = livyProperties;
        /*
        session retry 10 * 5000 = 50000
         */
        this.sessionRetryTemplate = retryTemplate(10, 5000L);
        /*
        statement retry 15 * 800 = 12000
         */
        this.statementRetryTemplate = retryTemplate(30, 800L);
        LivyProperties.LivyRetryConf livyRetryConf = livyProperties == null ? null : livyProperties.getLivyRetryConf();
        if (livyRetryConf != null) {
            if (livyRetryConf.getSessionMaxAttempts() > 0 && livyRetryConf.getSessionRetryTime() > 0) {
                this.sessionRetryTemplate = retryTemplate(livyRetryConf.getSessionMaxAttempts(),
                                livyRetryConf.getSessionRetryTime());
            }
            if (livyRetryConf.getStatementMaxAttempts() > 0 && livyRetryConf.getStatementRetryTime() > 0) {
                this.statementRetryTemplate = retryTemplate(livyRetryConf.getStatementMaxAttempts(),
                                livyRetryConf.getStatementRetryTime());
            }
        }
    }

    /**
     * 创建session
     *
     * @param kind
     * @return
     * @throws SessionException
     */
    public Session createSession(String kind) throws SessionException {
        try {
            LivySessionReq sessionReq = new LivySessionReq();
            sessionReq.setKind(kind);
            sessionReq.setName(buildName(kind));
            LivyProperties.SessionConf sessionConf = getLivyProperties().getSessionConf();
            if (sessionConf != null) {
                sessionReq.setQueue(sessionConf.getQueue());
                sessionReq.setConf(sessionConf.getConf());
            }
            if (sessionReq.getConf() == null) {
                sessionReq.setConf(new HashMap<>());
            }
            LivySessionRes sessionRes = livyCreateSession(sessionReq);
            Preconditions.checkState(sessionRes != null);
            Session session = new Session(this);
            session.setId(sessionRes.getId());
            return session;
        } catch (Exception e) {
            throw new SessionException("session create error", e);
        }
    }

    /**
     * 创建Statement，会重试几次直到成功或者超时
     *
     * @param sessionId
     * @param sql
     * @return
     * @throws SessionException
     */
    public Statement createStatement(Long sessionId, String sql) throws SessionException {
        return sessionRetryTemplate.execute(context -> createStatementInternal(sessionId, sql));
    }

    /**
     * 查询Statement，会重试几次直到成功或者超时
     *
     * @param sessionId
     * @param statementId
     * @return
     * @throws SessionException
     */
    public LivyStatementOutputQo.OutputDataJson queryStatement(Long sessionId, Long statementId)
                    throws SessionException {
        return statementRetryTemplate.execute(context -> queryStatementInternal(sessionId, statementId));
    }

    /**
     * 关闭session
     *
     * @param sessionId
     */
    public void killSession(Long sessionId) {
        try {
            livyKillSession(sessionId);
        } catch (HttpClientErrorException.NotFound e) {
            log.warn("kill session not found {}", e.getResponseBodyAsString());
        } catch (Exception e) {
            log.error("kill session error ignore", e);
        }
    }

    /**
     * 检查session 是否有效
     *
     * @param sessionId
     * @return
     */
    public boolean checkSession(Long sessionId) {
        try {
            livyCheckSession(sessionId);
            return true;
        } catch (HttpClientErrorException.NotFound e) {
            log.warn("check session not found {}", e.getResponseBodyAsString());
            return false;
        } catch (Exception e) {
            log.error("check session error ignore", e);
            return false;
        }
    }

    public List<Session> getActiveSession(SessionKind sessionKind) {
        try {
            String livyUrl = getSessionUrl();
            LivySessionsRes sessionsRes = restTemplate.getForObject(livyUrl, LivySessionsRes.class);
            Preconditions.checkState(sessionsRes != null);
            List<LivySessionRes> list = sessionsRes.getSessions();
            String prefix = buildNamePrefix(sessionKind.getName());
            if (list != null && !list.isEmpty()) {
                return list.stream().filter(s -> {
                    // 过滤name是同一个分组，根据name前缀
                    // 过滤状态为active的session
                    boolean active = s.getState().isActive();
                    return active && StringUtils.startsWith(s.getName(), prefix);
                }).map(s -> {
                    Session session = new Session(this);
                    session.setId(s.getId());
                    return session;
                }).collect(Collectors.toList());
            }
        } catch (Exception e) {
            log.error("get active session error ignore", e);
        }
        return Lists.newArrayList();
    }

    /**
     * livy 请求
     *
     * @param sessionReq
     * @return
     */
    public LivySessionRes livyCreateSession(LivySessionReq sessionReq) {
        String livyUrl = getSessionUrl();
        return restTemplate.postForObject(livyUrl, sessionReq, LivySessionRes.class);
    }

    public void livyKillSession(Long sessionId) {
        String livyUrl = getSessionIdUrl(sessionId);
        restTemplate.delete(livyUrl);
    }

    public void livyCheckSession(Long sessionId) {
        String livyUrl = getSessionIdUrl(sessionId);
        restTemplate.getForObject(livyUrl, LivySessionRes.class);
    }

    public LivyStatementRes livyCreateStatement(Long sessionId, LivyStatementReq statementReq) {
        String livyUrl = getStatementUrl(sessionId);
        return restTemplate.postForObject(livyUrl, statementReq, LivyStatementRes.class);
    }

    public LivyStatementRes livyQueryStatement(Long sessionId, Long statementId) {
        String livyUrl = getStatementIdUrl(sessionId, statementId);
        return restTemplate.getForObject(livyUrl, LivyStatementRes.class);
    }

    private Statement createStatementInternal(Long sessionId, String sql) throws SessionException {
        try {
            LivyStatementReq statementReq = new LivyStatementReq();
            statementReq.setCode(sql);
            statementReq.setKind(SessionKind.SQL.getName());

            LivyStatementRes statementRes = livyCreateStatement(sessionId, statementReq);
            Preconditions.checkState(statementRes != null);
            Statement statement = new Statement();
            statement.setId(statementRes.getId());
            if (log.isDebugEnabled()) {
                log.debug("LSSP {}-{} create sql: {}", sessionId, statement.getId(), sql);
            }
            return statement;
        } catch (HttpServerErrorException.InternalServerError e) {
            // retry
            log.warn("create statement 5xx, will retry, reason: {}", e.getResponseBodyAsString());
            throw new SessionNotReadyException("session not ready id = " + sessionId);
        } catch (Exception e) {
            throw new SessionException("statement create error", e);
        }
    }

    private LivyStatementOutputQo.OutputDataJson queryStatementInternal(Long sessionId, Long statementId)
                    throws SessionException {
        try {
            LivyStatementRes statementRes = livyQueryStatement(sessionId, statementId);
            StatementState state = null;
            if (statementRes != null) {
                state = statementRes.getState();
                LivyStatementOutputQo output = statementRes.getOutput();

                if (StatementState.WAITING == state || StatementState.RUNNING == state || output == null) {
                    throw new StatementNotReadyException("statement not ready id = " + statementId);
                }
                if (StatementState.AVAILABLE == state) {
                    OutputStatus status = output.getStatus();
                    if (log.isDebugEnabled()) {
                        log.debug("LSSP {}-{} query cost: {}ms, sql: {}", sessionId, statementId, getCost(statementRes),
                                        statementRes.getCode());
                    }
                    if (OutputStatus.OK == status) {
                        return output.getData().getDataJson();
                    } else if (OutputStatus.ERROR == status) {
                        SessionException sessionException = new SessionException(output.getEvalue());
                        StackTraceElement[] first = wrapStackTrace(output.getTraceback());
                        StackTraceElement[] second = sessionException.getStackTrace();
                        StackTraceElement[] all = ArrayUtils.addAll(first, second);
                        sessionException.setStackTrace(all);
                        throw sessionException;
                    }
                }
            }
            throw new SessionException("statement running error state is " + state);
        } catch (SessionException e) {
            throw e;
        } catch (Exception e) {
            throw new SessionException("statement running error", e);
        }
    }

    private StackTraceElement[] wrapStackTrace(List<String> traceback) {
        if (traceback == null) {
            return new StackTraceElement[0];
        }
        List<StackTraceElement> stacks = Lists.newArrayList();
        for (String trace : traceback) {
            try {
                Matcher matcher = PATTERN.matcher(trace);
                if (matcher.matches()) {
                    StackTraceElement element = new StackTraceElement(matcher.group(1), matcher.group(2),
                                    matcher.group(3), Integer.parseInt(matcher.group(4)));
                    stacks.add(element);
                }
            } catch (Exception e) {
                log.warn("wrap stack trace error {}", trace, e);
            }
        }
        return stacks.toArray(new StackTraceElement[0]);
    }

    private String getCost(LivyStatementRes statementRes) {
        if (statementRes.getCompleted() == null || statementRes.getStarted() == null) {
            return "-";
        }
        return String.valueOf(statementRes.getCompleted() - statementRes.getStarted());
    }

    private RetryTemplate retryTemplate(int maxAttempts, long retryTime) {
        Map<Class<? extends Throwable>, Boolean> retryableExceptions = new HashMap<>(16);
        retryableExceptions.put(SessionNotReadyException.class, true);
        retryableExceptions.put(StatementNotReadyException.class, true);
        SimpleRetryPolicy retryPolicy = new SimpleRetryPolicy(maxAttempts, retryableExceptions);

        FixedBackOffPolicy backOffPolicy = new FixedBackOffPolicy();
        backOffPolicy.setBackOffPeriod(retryTime);
        RetryTemplate retryTemplate = new RetryTemplate();
        retryTemplate.setRetryPolicy(retryPolicy);
        retryTemplate.setBackOffPolicy(backOffPolicy);
        return retryTemplate;
    }

    public LivyProperties getLivyProperties() {
        return livyProperties == null ? new LivyProperties() : livyProperties;
    }

    private String buildName(String kind) {
        String prefix = buildNamePrefixAddIp(kind);
        return prefix + System.currentTimeMillis();
    }

    private String buildNamePrefix(String kind) {
        return String.format(SESSION_PREFIX, kind.toUpperCase());
    }

    private String buildNamePrefixAddIp(String kind) {
        return buildNamePrefix(kind) + NetUtil.getLocalhostStr() + "-";
    }

    private String getSessionUrl() {
        return String.format(SESSION_URL, livyHost);
    }

    private String getSessionIdUrl(Long sessionId) {
        return String.format(SESSION_ID_URL, livyHost, sessionId);
    }

    private String getStatementUrl(Long sessionId) {
        return String.format(STATEMENT_URL, livyHost, sessionId);
    }

    private String getStatementIdUrl(Long sessionId, Long statementId) {
        return String.format(STATEMENT_ID_URL, livyHost, sessionId, statementId);
    }
}
